import tensorflow as tf


def LSTM_Model(model, input_data, output_data, vocab_size
               , run_size=128, num_layers=2, batch_size=64, learning_rate=0.01):

    end_points = {}
    if model == 'rnn':
        cell_fun = tf.contrib.rnn.BasicRNNCell
    elif model == 'gru':
        cell_fun = tf.contrib.rnn.GRUCell
    elif model == 'lstm':
        cell_fun = tf.contrib.rnn.BasicLSTMCell    # rnn基础模型
    cell = cell_fun(run_size, state_is_tuple=True)
    cell = tf.contrib.rnn.MultiRNNCell([cell]*num_layers, state_is_tuple=True)

    if output_data is not None:
        initial_state = cell.zero_state(batch_size, tf.float32)
    else:
        initial_state = cell.zero_state(1, tf.float32)

    with tf.device('/cpu:0'):
        embedding = tf.get_variable('embedding', initializer=tf.random_uniform([vocab_size+1, run_size], -1.0, 1.0))
        inputs = tf.nn.embedding_lookup(embedding, input_data)
    outputs, last_state = tf.nn.dynamic_rnn(cell, inputs, initial_state=initial_state)
    output = tf.reshape(outputs, [-1, run_size])

    weights = tf.Variable(tf.truncated_normal([run_size, vocab_size+1]))
    bias = tf.Variable(tf.zeros(shape=[vocab_size+1]))
    logits = tf.nn.bias_add(tf.matmul(output, weights), bias=bias)

    if output_data is not None:
        labels = tf.one_hot(tf.reshape(output_data, [-1]), depth=vocab_size+1)
        loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)
        total_loss = tf.reduce_mean(loss)   # 求平均
        train_op = tf.train.AdamOptimizer(learning_rate).minimize(total_loss)

        end_points['initial_state'] = initial_state
        end_points['output'] = output
        end_points['train_op'] = train_op
        end_points['total_loss'] = total_loss
        end_points['loss'] = loss
        end_points['last_state'] = last_state
    else:
        prediction = tf.nn.softmax(logits)

        end_points['initial_state'] = initial_state
        end_points['last_state'] = last_state
        end_points['prediction'] = prediction

    return end_points



